import cv2
import numpy as np

def center(P_left_top,P_right_top,P_left_bot,P_right_bot):
    w=P_right_top[0]-P_left_top[0]
    h=P_left_bot[1]-P_left_top[1]
    x0=P_left_top[0]+w/2
    y0=P_left_top[1]+h/2
    return np.array([[np.float32(x0)],[np.float32(y0)]])    #返加卡尔曼滤波使用的列向量



class Pedestrain():
    """Pedestrain类,由ROI，ID,卡尔曼滤波器组成，用Pedestrain来保持目标的状态"""

    def __init__(self,id,frame, track_window):
        ##用追踪窗口坐标来初始化行人目标
        self.id=int(id)
        x,y,w,h=track_window
        self.track_window=track_window
        self.roi = cv2.cvtColor(frame[y:y+h,x:x+w],cv2.COLOR_BGR2HSV)
        roi_hist = cv2.calcHist([self.roi], [0], None, [16],[0,180])
        self.roi_hist = cv2.normalize(roi_hist,roi_hist,0,255,cv2.NORM_MINMAX)
        ##设置Kalman滤波
        self.kalman=cv2.KalmanFilter(4,2)       #状态变量4个维度，两个测量数据2个维度
        #观测矩阵
        self.kalman.measurementMatrix= np.array([[1,0,0,0],[0,0,0,0]],np.float32)
        #状态转移矩阵
        self.kalman.transitionMatrix=np.array([[1,0,1,0],[0,1,0,1],[0,0,1,0],[0,0,0,1]],np.float32)
        #噪声协方差矩阵
        self.kalman.processNoiseCov=np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], np.float32)*0.03
        self.measurement=np.array((2,1),np.float32)
        self.prediction = np.zeros((2,1),np.float32)
        self.term_crit=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10,1)
        self.center=None
        self.font=cv2.FONT_HERSHEY_SIMPLEX
        self.update(frame)


    def __del__(self):
        print("Pedestrain %d destroyed",self.id)

    def update(self,frame,**kwargs):
        #对frame进行反投影
        hsv=cv2.cvtColor(frame,cv2.COLOR_BGR2HSV)
        back_project = cv2.calcBackProject([hsv],[0],self.roi_hist,[0,180],1)

        if kwargs.get("algorithm") == "c":
            ret,self.track_window=cv2.CamShift(back_project, self.track_window,self.term_crit)
            pts=cv2.boxPoints(ret)
            pts=np.int0(pts)
            self.center = center(pts)
            cv2.polylines(frame,[pts],True, 255,1)
        if not kwargs.get("algorithm") or kwargs.get("algorithm") == "m" :
            ret,self.track_window = cv2.meanShift(back_project, self.track_window,self.term_crit)
            x,y,w,h=self.track_window
            self.center =center([x,y],[x+w,y],[x,y+h],[x+w,y+h])
            cv2.rectangle(frame,(x,y),(x+w,y+h),(255,255,0),1)

        self.kalman.correct(self.center)    #从测量数据更新预测状态
        prediction = self.kalman.predict()  #
        cv2.circle(frame,(int(prediction[0]),int(prediction[1])),4, (0,255,0), -1)

        cv2.putText(frame,"ID:{0}->{1}".format(self.id,self.center),(11,(self.id+1)*25+1),
                    self.font,0.6, (0,255,0),1, cv2.LINE_AA)
        cv2.putText(frame,"ID:{0}->{1}".format(self.id,self.center),(10,(self.id+1)*25),
                    self.font,0.6, (0,255,0),1, cv2.LINE_AA)


def main():
    #camera = cv2.VideoCapture();
    camera = cv2.VideoCapture("../video/room.avi")
    #camera = cv2.VideoCapture(0);

    """首先载入视频，然后初始化背景提取器，设置20帧作为影响背景模型的帧数"""
    history =20
    bs= cv2.createBackgroundSubtractorKNN(detectShadows=True)
    bs.setHistory(history)
    """创建主显示窗口，然后建立一个行人字典，firstFrame标志"""
    cv2.namedWindow("surveillance")
    pedestrians={}
    firstFrame =True
    frames = 0

    ##进入工作循环，从camera一帧一帧的读取图像
    while True:
        print("------------------FRAME %d ------------------",frames);
        grabbed, frame = camera.read()
        if(grabbed is False):
            print("failed to grab frame")
            break
        ret, frame = camera.read()
        """先让BackgroundSubstractorKNN对背景模型创建历史，所以我们实际不处理前20帧，
        而是交给背景提取器"""
        fgmask = bs.apply(frame)
        if frames<history:
            frames+=1
            continue

        """获取轮廓"""
        th = cv2.threshold(fgmask.copy(), 127,255,cv2.THRESH_BINARY)[1]
        th = cv2.erode(th,cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3)),iterations=2)
        dilated =cv2.dilate(th, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(8,3)),iterations=2)
        image, contours, hier=cv2.findContours(dilated,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
        """一旦轮廓被确认，我们仅为第一帧(除开前面的history帧数)的每个轮廓实例化一个Pedestrain类对象"""
        counter =0
        for c in contours:
            area=cv2.contourArea(c)
            if area> 500:
                (x,y,w,h)=cv2.boundingRect(c)
                cv2.rectangle(frame,(x,y),(x+w,y+h),(0,255,0),1)
                if firstFrame is True:
                    pedestrians[counter]=Pedestrain(counter,frame,(x,y,w,h))
                counter +=1
        """然后对于每一个检测到的行人，执行update方法（传递当前帧做参数）"""
        for i,p in pedestrians.items():
            p.update(frame)
        """设置firstFrame为False,这样就不会实例化更多的pedestrain类了"""
        firstFrame =False
        frames+=1

        cv2.imshow("surveillance",frame)
        if cv2.waitKey(110) & 0xff==27:
            break

#当文件作为模块时，不执行main()
if __name__ == "__main__":
    main()
